import os
import shutil
from datetime import datetime
import cv2

class DatasetInitializer:
    def __init__(self, dataset_dir='dataset'):
        self.dataset_dir = dataset_dir
        self.target_dir = os.path.join(dataset_dir, 'target')
        self.background_dir = os.path.join(dataset_dir, 'background')
        
    def initialize_folders(self):
        """初始化数据集文件夹结构"""
        print("\n开始初始化数据集文件夹...")
        
        # 创建主文件夹
        os.makedirs(self.dataset_dir, exist_ok=True)
        os.makedirs(self.target_dir, exist_ok=True)
        os.makedirs(self.background_dir, exist_ok=True)
        
        print("✓ 创建文件夹结构完成")
        print(f"- 目标物品文件夹: {self.target_dir}")
        print(f"- 背景图片文件夹: {self.background_dir}")
    
    def rename_images(self):
        """重命名所有图片"""
        print("\n开始重命名图片...")
        
        # 重命名目标物品图片
        self._rename_folder_images(self.target_dir, 'target')
        
        # 重命名背景图片
        self._rename_folder_images(self.background_dir, 'background')
        
    def _rename_folder_images(self, folder_path, prefix):
        """重命名指定文件夹中的图片"""
        if not os.path.exists(folder_path):
            print(f"文件夹不存在: {folder_path}")
            return
            
        # 获取所有图片文件
        image_files = [f for f in os.listdir(folder_path) 
                      if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
        
        if not image_files:
            print(f"警告: {folder_path} 中没有找到图片")
            return
            
        print(f"\n处理 {prefix} 文件夹:")
        print(f"找到 {len(image_files)} 张图片")
        
        # 重命名图片
        for i, old_name in enumerate(image_files, 1):
            # 获取文件扩展名
            _, ext = os.path.splitext(old_name)
            
            # 生成新文件名 (例如: target_001.jpg)
            new_name = f"{prefix}_{i:03d}{ext.lower()}"
            
            # 完整的文件路径
            old_path = os.path.join(folder_path, old_name)
            new_path = os.path.join(folder_path, new_name)
            
            try:
                # 检查图片是否可以打开
                img = cv2.imread(old_path)
                if img is None:
                    print(f"警告: 无法读取图片 {old_name}")
                    continue
                
                # 重命名文件
                os.rename(old_path, new_path)
                print(f"✓ {old_name} -> {new_name}")
                
            except Exception as e:
                print(f"错误: 处理 {old_name} 时出错: {str(e)}")
    
    def verify_dataset(self):
        """验证数据集"""
        print("\n验证数据集...")
        
        # 检查目标物品图片
        target_images = [f for f in os.listdir(self.target_dir) 
                        if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
        
        # 检查背景图片
        background_images = [f for f in os.listdir(self.background_dir) 
                           if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
        
        print("\n数据集统计:")
        print(f"- 目标物品图片: {len(target_images)} 张")
        print(f"- 背景图片: {len(background_images)} 张")
        print(f"- 总计: {len(target_images) + len(background_images)} 张")
        
        if len(target_images) < 50:
            print("\n⚠️ 警告: 目标物品图片数量不足，建议至少50张")
        if len(background_images) < 50:
            print("\n⚠️ 警告: 背景图���数量不足，建议至少50张")

def main():
    """主函数"""
    initializer = DatasetInitializer()
    
    # 初始化文件夹
    initializer.initialize_folders()
    
    # 等待用户放入图片
    input("\n请将图片放入对应文件夹后，按回车键继续...")
    
    # 重命名图片
    initializer.rename_images()
    
    # 验证数据集
    initializer.verify_dataset()
    
    print("\n数据集初始化完成！")
    print("\n下一步:")
    print("1. 确保每个文件夹都有足够的图片（建议每类至少50张）")
    print("2. 运行 python src/main.py 开始训练模型")

if __name__ == "__main__":
    main() 